#include "StdAfx.h"
#include "WMIHandler.h"

CWMIHandler::CWMIHandler(void)
: m_pWbemLocator(NULL)
, m_pWbemSvc(NULL)
, m_bConnected(false)
, m_bNeedUninitCom(true)
{
}

CWMIHandler::~CWMIHandler(void)
{
  cleanup();
  CoUninitialize();
}

// If the COM has been initialized, the function will return false
bool CWMIHandler::initCOMAndSecurity(DWORD dwCoInit)
{
  HRESULT hRes;

  hRes = CoInitializeEx(0, dwCoInit);
  if ( S_OK != hRes
      && S_FALSE != hRes )   // if the COM has been initialized, it will return S_FALSE
  {
    // CoInitializeEx() failed
    return false;
  }
  else if ( S_FALSE == hRes )
  {
    m_bNeedUninitCom = false;
  }

  hRes = CoInitializeSecurity(
    NULL,
    -1,
    NULL,
    NULL,
    RPC_C_AUTHN_LEVEL_DEFAULT,
    RPC_C_IMP_LEVEL_IMPERSONATE,
    NULL,
    EOAC_NONE,
    NULL
    );

  if ( FAILED(hRes) )
  {
    if (m_bNeedUninitCom)
    {
      CoUninitialize();
    }
    else
    {
      m_bNeedUninitCom = true;
    }
    return false;
  }

	return true;
}

ULONG CWMIHandler::connectToNamespace(LPCTSTR lpcszWmiNamespace)
{
  HRESULT hRes;

  hRes = CoCreateInstance(
    CLSID_WbemLocator,
    0,
    CLSCTX_INPROC_SERVER,
    IID_IWbemLocator, (LPVOID *)&m_pWbemLocator);
  if ( FAILED(hRes) )
  {
    cleanup();
    return WMI_HANDLER_RTCODE_CREATE_LOCATOR_FAILED;
  }

  hRes = m_pWbemLocator->ConnectServer(
      _bstr_t(lpcszWmiNamespace),
      NULL,
      NULL,
      0,
      NULL,
      0,
      0,
      &m_pWbemSvc);
  if ( FAILED(hRes) )
  {
    cleanup();
    return WMI_HANDLER_RTCODE_CONNECT_NAMESPACE_FAILED;
  }

  hRes = CoSetProxyBlanket(
    m_pWbemSvc,
    RPC_C_AUTHN_WINNT,
    RPC_C_AUTHZ_NONE,
    NULL,
    RPC_C_AUTHN_LEVEL_CALL,
    RPC_C_IMP_LEVEL_IMPERSONATE,
    NULL,
    EOAC_NONE);
  if ( FAILED(hRes) )
  {
    cleanup();
    return WMI_HANDLER_RTCODE_SET_PROXY_FAILED;
  }

  m_bConnected = true;

	return WMI_HANDLER_RTCODE_OK;
}

ULONG CWMIHandler::registerEventSink(LPCTSTR lpcszQuery, CWMIEventSink * pEventSink)
{
  if ( !isConnected() )
  {
    return WMI_HANDLER_RTCODE_NOT_CONNECTED;
  }

  if ( pEventSink->isRegistered() )
  {
    return WMI_HANDLER_RTCODE_EVENTSINK_HAS_BEEN_REGISTERED;
  }

  return pEventSink->registerSelf(lpcszQuery, m_pWbemSvc);
}

ULONG CWMIHandler::unregisterEventSink(CWMIEventSink * pEventSink)
{
  if ( !isConnected() )
  {
    return WMI_HANDLER_RTCODE_NOT_CONNECTED;
  }

  if ( !pEventSink->isRegistered() )
  {
    return WMI_HANDLER_RTCODE_OK;
  }

  return pEventSink->unregisterSelf();
}

void CWMIHandler::cleanup(void)
{
  if ( m_pWbemSvc != NULL )
  {
    m_pWbemSvc->Release();
    m_pWbemSvc = NULL;
  }

  if ( m_pWbemLocator != NULL )
  {
    m_pWbemLocator->Release();
    m_pWbemLocator = NULL;
  }

  m_bConnected = false;
}


ULONG CWMIHandler::execQuery(LPCTSTR lpcszQuery, IEnumWbemClassObject ** ppEnumerator)
{
  HRESULT hRes;

  if ( !isConnected() )
  {
    return WMI_HANDLER_RTCODE_NOT_CONNECTED;
  }

  hRes = m_pWbemSvc->ExecQuery(
    TEXT("WQL"),
    _bstr_t(lpcszQuery),
    WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
    NULL,
    ppEnumerator);
  if ( FAILED(hRes) )
  {
    return WMI_HANDLER_RTCODE_EXEC_QUERY_FAILED;
  }

  return WMI_HANDLER_RTCODE_OK;
}
